package persistence

import (
	"context"

	"Data-API-MCP/backend/internal/domain/entity"
	"Data-API-MCP/backend/internal/domain/repository"
	"sync"
)

// ConnectionManager 数据库连接管理器
type ConnectionManager struct {
	connections sync.Map
}

// NewConnectionManager 创建连接管理器
func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{}
}

// GetConnection 获取数据库连接
func (m *ConnectionManager) GetConnection(ctx context.Context, conn *entity.Connection) (repository.DBConnection, error) {
	// 检查连接是否已存在
	if existingConn, ok := m.connections.Load(conn.ID); ok {
		return existingConn.(repository.DBConnection), nil
	}

	// 创建新连接
	dbConn, err := m.createConnection(ctx, conn)
	if err != nil {
		return nil, err
	}

	// 存储连接
	m.connections.Store(conn.ID, dbConn)
	return dbConn, nil
}

// CloseConnection 关闭数据库连接
func (m *ConnectionManager) CloseConnection(ctx context.Context, conn *entity.Connection) error {
	if existingConn, ok := m.connections.Load(conn.ID); ok {
		dbConn := existingConn.(repository.DBConnection)
		if err := dbConn.Close(); err != nil {
			return err
		}
		m.connections.Delete(conn.ID)
	}
	return nil
}

// CloseAllConnections 关闭所有连接
func (m *ConnectionManager) CloseAllConnections() error {
	m.connections.Range(func(key, value interface{}) bool {
		if dbConn, ok := value.(repository.DBConnection); ok {
			dbConn.Close()
		}
		m.connections.Delete(key)
		return true
	})
	return nil
}

// createConnection 创建新的数据库连接
func (m *ConnectionManager) createConnection(ctx context.Context, conn *entity.Connection) (repository.DBConnection, error) {
	// 根据连接类型创建不同的数据库连接
	switch conn.Type {
	case "mysql":
		return NewMySQLConnection(conn)
	case "postgresql":
		return NewPostgreSQLConnection(conn)
	default:
		return nil, repository.ErrUnsupportedDatabaseType
	}
}
